  
open Option
open ExtList
open Printf
open ExtLib


open Constants
open Management
open Types
open Randif
open Bmap
open Banking;;


type cash_rate_params = float * float * float

type regulator = float * float

(*
* banks interacting in a world controlled by management
*
* This module describes transactions and overall world-stepping
*)

type world = { wdate: date; (* world-date - ie. current date *)
	       enddate: date;
               cash_rate_params: cash_rate_params;
               banks: bank list;
               lenders: contractFactory list;
               borrowers: contractFactory list;
               regulator: regulator;
	       diagnostic: unit }


(* mean, amplitude, period *)
let cash_rate (m,a,p) d = 
  m +. a *. (sin ( 2. *. pi *. (dateFloat d) /. p))

let cash_rate_world w =
  cash_rate w.cash_rate_params w.wdate


(*
newProposal :: Gen a -> Date -> ContractFactory -> ST a BondProposal
-- newContract d g f | trace ("newContract : " ++ (show f)) False = undefined   
*)
let new_proposal d { defaultDist = dd; principalDist = pd; rateStart = rstart; 
		    durationDist = durd; recoveryDist = recd }
    = let gdur = (getDNum durd) in
      (* let grstep = (getDNum rstep) in  *)
      let grstart = (getDNum rstart) in
      let gprin = (getDNum pd) in   
      let gdef = (getDNum dd) in
      let grec = (getDNum recd) in
      {start = d; maturity = sumDate d (date_of_float gdur);
       rate = grstart; principal = gprin ;rdefault = gdef; recovery = grec}

(*
* adjustProposal :: Gen a -> ContractFactory -> BondProposal -> ST a BondProposal
*)

(*

let adjustProposal {rateStep = rs} proposal
    = let adjustment = (getDNum rs) in
      let { rate = rt } = proposal in
      { proposal with rate = rt +. adjustment }

*)

(*
evaluateContractP :: Date -> Bank -> Policy -> BondProposal -> Double -> Int -> Int -> (Bool,Int) -- (result,cost)
*)


let rec evaluateContractP pol d cr proposal bank instrs maxinstrs =
  if instrs < maxinstrs
  then evaluateContractP' pol d cr proposal bank (instrs + 1) maxinstrs
  else (false,instrs)
and evaluateContractP' pol d cr proposal bank instrs maxinstrs =
  (* let mgt = bank.management in *)
  (* let m = match mgt with
          | (SimpleM x) -> (SimpleM x)
          | _ -> NullM *)
  match pol with 
  | PAccept -> (true,instrs)
  | PReject -> (false,instrs)
  | (PCall n) -> 
      (match bank.management with
      | (SimpleM m) ->
	  (evaluateContractP (lkupRegister m n) d cr proposal bank instrs maxinstrs)
      | _ -> raise INTERNAL_ERROR )
  | (PIf (test,r1,r2)) ->
      let (res,instrs0) = evaluateContractP pol d cr proposal bank instrs maxinstrs in
      if res
      then evaluateContractP r1 d cr proposal bank instrs0 maxinstrs 
      else evaluateContractP r2 d cr proposal bank instrs0 maxinstrs 
  | (PNot p) ->
      let (res,instrs0) = evaluateContractP pol d cr proposal bank instrs maxinstrs in
      (not res,instrs0)
  | (PNumTest (n1,dir,n2)) ->
      let (n1res,instrs0) = (evalNumP n1 d cr proposal bank instrs maxinstrs) in
      let (n2res,instrs1) = (evalNumP n2 d cr proposal bank instrs0 maxinstrs) in
      ((lookupDir dir) n1res n2res, instrs1)
  | PNullPolicy -> raise INTERNAL_ERROR;

and evalNumP pnum d cr proposal b i maxi =
  let  { bankcash = bc } = b in
  let  { start = st;
	 maturity = mat;
	 rate = r;
	 principal = p;
	 rdefault = rd;
	 recovery = reco} = proposal in
  match pnum with
  | PRate ->  (r*.rateScale,i+1)
  | PPrincipal ->  (p*.principalScale,i+1)
  | PDefault ->  (rd*.defScale,i+1)
  | PMaturity ->
      (matScale *. (dateFloat ( sumDate mat d)),i+1)
  | PBankCash ->  (bcScale *. bc,i+1)
  | PCapitalRat -> (b.capital_ratio,i+1)
  | PLiqRat -> (b.liquid_ratio,i+1)
  | (PNum n) -> (n,i+1)
  | PCashRate -> (crScale *. cr,i+1)
  | (PSum (n1,n2)) -> 
    let (s1,i0) = (evalNumP n1 d cr proposal b i maxi) in
    let (s2,i1) = (evalNumP n1 d cr proposal b i0 maxi) in
    (s1 +. s2,i1)
  | (PDiff (n1,n2)) -> 
    let (s1,i0) = (evalNumP n1 d cr proposal b i maxi) in
    let (s2,i1) = (evalNumP n1 d cr proposal b i0 maxi) in
    (s1 -. s2,i1)
  | (PProd (n1,n2)) ->
    let (s1,i0) = (evalNumP n1 d cr proposal b i maxi) in
    let (s2,i1) = (evalNumP n1 d cr proposal b i0 maxi) in
    (s1 *. s2,i1)
  | (PDiv (n1,n2)) ->
    let (s1,i0) = (evalNumP n1 d cr proposal b i maxi) in
    let (s2,i1) = (evalNumP n1 d cr proposal b i0 maxi) in
    (s1 /. s2,i1)
  | (PBSGap (n1,n2)) ->
      (match (n1,n2) with 
      | (PNum m1,PNum m2) ->
	  let gap = bsgap d b m1 m2 in
	  (gap,i+1)
      | (_,_) -> raise INTERNAL_ERROR)
  | (PBSSum (_, _)|PNullNum) -> raise INTERNAL_ERROR
  (* | _ -> raise INTERNAL_ERROR *)



let lookup_index rb r = 
  let rec li1 (rb : float list) ( r : float) ( n : int ) =
    match rb with
    | [] -> n
    | b :: _ when r <= b -> n
    | b :: bs when r > b   -> li1 bs r (n + 1) 
    | _ -> raise INTERNAL_ERROR  
  in
  li1 rb r 0

		      
(*
* experience-based management
*)

(*
* - skeleton for testing/calibration

Ultimately this will be the product of the different individual
posterior odds ...

*)

(*
let acceptance_probability party cashrate bank proposal = 0.5
*)

let construct_pd crw d b c =
  let cstemp = contract_state_build crw d b c false in
 (List.drop cs_final cstemp.cs_state)

let getptybm b party =
  match (b.management,party) with
  | (BayesM m,Borrower) -> (m.accepted_borrowers,m.rejected_borrowers)
  | (BayesM m,Lender) -> (m.accepted_lenders,m.rejected_lenders)
  | (_,_) -> print_string "ArRK!\n"; 
      raise INTERNAL_ERROR

let acceptance_probability party crw b c =
  let (accbm,rejbm) = getptybm b party in
  let prior_data = construct_pd crw zeroDate b c in
  let probsacc = List.map2 (hiprof accbm) (cs_final -- (cs_max - 1)) prior_data in
  let probsrej = List.map2 (hiprof rejbm) (cs_final -- (cs_max - 1)) prior_data in
  let prod = List.fold_left ( *. ) 1.0 in
  let pa = prod probsacc in
  let pr = prod probsrej in
  (* printf "pa is %f pr is %f final %f\n%!" pa pr (pa /. (pa +. pr) ); *)
  pa /. (pa +. pr) 


let evaluateProposal dt cashrate pty proposal bank = 
  (* if bank.bankID == 1
  then Printf.printf "BankID 1 - proposal %s party %s ..." (Std.dump proposal) (print_party pty)
  else (); *)
  let (testres,count) = match bank.management with
                        | (SimpleM mgt) -> 
			    let pol = match pty with
			    | Lender -> mgt.lenderPolicy
			    | Borrower -> mgt.borrowerPolicy in
			    (evaluateContractP pol dt cashrate proposal bank 0 maxInstructions)
(*
                        | (TableM mgt) -> 
			    (let tbl = match pty with
			    | Lender -> (mgt.lender_table,mgt.lender_tests)
			    | Borrower -> (mgt.borrower_table,mgt.borrower_tests) in
			    let restp = (evaluateTableP tbl dt cashrate proposal bank mgt) in
			    let rn = Random.float 1.0 in
			    (match restp with
			    | true -> (* (rn > 0.2,0) *) (true,0)
			    | false -> (* (rn > 0.8,0) *) (false,0)))
*)
			| NullM -> raise INTERNAL_ERROR
  in
  (* if bank.bankID == 1
  then 
    Printf.printf "%s"
    (match testres with
     | true -> " accepted\n"
     | false -> " rejected\n")
  else (); *)
  if testres
  then (true,accept_contract pty dt bank proposal)
  else (false,reject_contract pty dt bank proposal)


  (* should make these argument order consistent *)

let rec amut f ls0 list =
  match list with 
  | [] -> (None,ls0)
  | (l :: ls) -> 
      let (res,b) = f l in
      if (res) 
      then (Some () , (List.rev ls0) @ (b :: ls))
      else amut f (b :: ls0) ls

let applyMapUntilTrue func list = amut func [] list

let checkProposalWithAllBanks dt prop cashrate pty banks =
  applyMapUntilTrue (evaluateProposal dt cashrate pty prop) banks

let rec runNUntil func x result count = 
  match result with 
  | None when count > 0 ->
      let (nresult,xp) = func x count in 
      runNUntil func xp nresult (count - 1)
  | None -> (None,x)
  | Some _ -> (Some (),x) 

(* should be bool - but we can keep the some/none for now *)

(*
let checkProposal dt cashrate bs factory =
  let { offeror = pty ; maxOfferSteps = times } = factory in
  let adjustAndCheckBanks (prop,banks) n =
    let nprop = adjustProposal factory prop in
    let (pstate,bks) = checkProposalWithAllBanks dt nprop cashrate pty banks in
    (pstate,(nprop,bks))
  in 
  let (res,(prop,nbs)) = runNUntil (adjustAndCheckBanks) (proposal,bs) None times in
  nbs (* eventually use res and prop for diagnostics *)
*)

(*
  let ap1 p cr c b = 

  let fst3 (a,_,_) = a in

  let aprobs = List.map (ap1 party cashrate proposal) bs in
  let saprobs = List.sort ~cmp:(fun a b -> compare (fst3 a) (fst3 b)) aprobs in
  print_endline ("aprobs = " ^ (Std.dump aprobs));

*)

let is_nan f = match classify_float f with FP_nan -> true | _ -> false

let rec test_proposal c cashrate date party maxtries = function 
  | [] -> []
  | (b:: bs) ->
      match maxtries with
      | 0 -> ((reject_contract party date b c)::(bs))
      | _ -> 
	  let ap = if b.regulator_control
	           then (0.0)
	           else (acceptance_probability party cashrate b c) in
	  let rnum = Random.float 1.0 in
	  (* printf "rnum %f ap %f party %s \n%!" rnum ap (Std.dump party); *)
	  (* if is_nan(ap)
	  then printf "rnum %f ap %f proposal %s \n%!" rnum ap (Std.dump c)
	  else ();  *)
	  if (rnum <= ap) 
	  then ((accept_contract party date b c)::(bs))
	  else ((reject_contract party date b c)::(test_proposal c cashrate date party (maxtries - 1) bs))

let check_proposal date cashrate bs offer =
  let proposal = (new_proposal date offer) in
  let party = offer.offeror in
  let nbs = test_proposal proposal cashrate date party max_proposal_tries bs in
  nbs

(* iterative transaction; checks each proposal against each bank *)

let rec transact dt cashrate bs offers =
  match offers with
  | [] -> bs
  | (o :: os) -> 
      let newbs0 = check_proposal dt cashrate bs o 
      in
      let rotbs = (rotate (List.rev newbs0)) in
      (*Printf.printf "Here! - headbank orig %s new %s rot %s offers %s\n%!" 
	(Std.dump (List.map (fun b -> b.bankID) bs))  
	(Std.dump (List.map (fun b -> b.bankID) newbs0))  
	(Std.dump (List.map (fun b -> b.bankID) rotbs))
	(Std.dump (List.length (o::os))); *)
      transact dt cashrate rotbs os


let regulate dt end_date (min_liq_ratio,min_capital_requirement) b 
    = if ((b.liquid_ratio > (min_liq_ratio)) &&
            (b.capital_ratio > (min_capital_requirement))) || (b.regulator_control)
      then b
      else {b with regulator_control = true; failTime = dateFloat (diffDate end_date dt)}


(*
*  1) transact
*  2) end of month accounts and interest payments
*  3) send reports to regulator who (instantly) decides if anyone gets the chop
*)


let step_world debug w =
  let { banks = bs; 
	wdate = dt;
	enddate = ed;
	borrowers = br;
	lenders = dep;
	regulator = reg;
	cash_rate_params = crp } = w in
  if debug then (print_endline "Calculating cash rate ...") else ();
  let cr = cash_rate crp dt in
  if debug then (print_endline "Building offers ...") else ();
  let allOffers = shuffle (br @ dep) in
  if debug then (print_endline "Transacting ...") else ();
  let bs1 = transact dt cr bs allOffers in
  if debug then (print_endline "Stepping time ...") else ();
  let bs2 = List.map (stepBank dt cr false) bs1 in
  if debug then (print_endline "Regulator regulamating ...") else ();
  let bs3 = List.map (regulate dt ed reg) bs2 in
  if debug then (print_endline "Constructing new world ...") else ();
  { w with banks = (shuffle bs3);
           wdate = (incDate dt) }


let step_world_done w =
  let { wdate = dt;
	enddate = ed } = w in
  (dateFloat ed) <= (dateFloat dt) 


(*
nStepWorld :: Gen a -> Bool -> Int -> World -> ST a World
nStepWorld g ka i w = foldM (\ (!x) _ -> stepWorld g ka x) w (replicate i 0)
*)

(*
{- briefsummary -}
bsummary :: Int -> SimState -> String
*)

let getBestBanks n w =
  let banks = w.banks in
  List.take n (List.sort ~cmp:(fun y x -> compare (sumDividends x) (sumDividends y)) banks)

let instances x t = List.length (List.filter (fun y -> y == x) t) 

let bsummary n c w 
    = (* let bestbanks = (getBestBanks 1 w) in *)
      let allbanks = w.banks in
      let solvbanks = (List.sort ~cmp:(fun y x -> compare (sumDividends x) (sumDividends y))
			 ( List.filter (fun x -> not (x.regulator_control)) allbanks)) in 
      let avggdivs = meanv ( List.map sumDividends (  solvbanks)) in
      (* betternew = [ b |  b <- newbanks, (sumDividends b) > avggdivs ]*)
      let avgfail = (meanv (List.map (fun x -> x.failTime) allbanks)) in
          (* avgnim = 100 * meanv (List.map (arithmean nim) solvbanks)
          avgroe = 100 * meanv (List.map (arithmean roe) solvbanks)
          avgroa = 100 * meanv (List.map (arithmean roa) solvbanks) *)
      let avgassets = meanv (List.map (fun b -> (valueBook b.assets) +. b.bankcash) solvbanks) in
      let avgliab = meanv (List.map (fun b -> (valueBook b.liabilities) +. b.deposits) solvbanks) in
      let avgequity = meanv (List.map (fun b -> b.book_equity) solvbanks) in
      let avgcap = 100. *. (meanv (List.map (fun b -> b.book_equity) solvbanks))/.(meanv (List.map riskWeightedAssets solvbanks)) in
      (* let insolvencies = instances true (List.map (fun x -> x.regulator_control) allbanks) in *)
      (* let insolvnew =  instances true (List.map (fun x -> x.regulator_control) (newbanks)) in  *)
      Printf.sprintf "%3d %5d %2d %2d %8.4f %15.4f %15.4f %15.4f %15.4f %15.4f\n"
	n c
	(List.length allbanks) (List.length solvbanks)
	avgfail
	avgassets
	avgliab
	avgequity
	avggdivs
	avgcap

let bsummary_header = " id  step tb so   avgttf       avgassets         avgliab       avgequity         avgdivs     avgcapratio\n"

	
let assignManagement mgts w =
  { w with banks = 
       List.map2 (fun b m -> {b with management = m }) 
           w.banks mgts }

let step_world_finalize w = 
  { w with banks = List.map bank_finalize w.banks }

let rec n_step_world debug w =
  if debug
  then 
    (
     Printf.printf "\n***\nWorld summary at date %f\n***\n" (dateFloat (w.wdate));
     print_endline (bsummary 0 0 w);
     print_endline (print_bank (List.hd (List.filter (fun b -> b.bankID == 1) w.banks))))
  else ();
  if step_world_done w 
  then (step_world_finalize w) (* finalize the state of the banks *)
  else n_step_world debug (step_world debug w)
